<!doctype html>
<html>
<head>
  <meta charset="utf-8">
  <meta name="viewport" content="width=device-width,initial-scale=1">
  <title>Edge AI Dual Models</title>
  <script src="js/tf.min.js"></script>
  <style>
    body {
      font-family: Arial, Helvetica, sans-serif;
      background: #181818;
      color: #EFEFEF;
      font-size: 16px;
    }
    #content {
      display: flex;
      flex-wrap: wrap;
      gap: 20px;
    }
    .panel {
      flex: 1;
      min-width: 100px;
      max-width: 210px;
      background: #363636;
      padding: 10px;
      border-radius: 8px;
    }
    textarea {
      width: 200px;
      height: 180px;
      resize: none;
      background: #111;
      color: #0f0;
      padding: 5px;
      border-radius: 4px;
      border: none;
    }
    img {
      max-width: 100%;
      border-radius: 8px;
      margin-top: 10px;
    }
  </style>
</head>
<body>
  <h2>本地双模型推理 Demo</h2>

  <!-- 图片上传 -->
  <input type="file" id="imageUpload" accept="image/*">
  <br>
  <img id="inputImage" width="300" alt="选择的图片">

  <div id="content">
    <div class="panel">
      <h3>模型 35-96</h3>
      <textarea id="result96" readonly></textarea>
    </div>
    <div class="panel">
      <h3>模型 100-224</h3>
      <textarea id="result224" readonly></textarea>
    </div>
  </div>

  <script>
    let model96, model224;

    async function loadModels() {
      try {
        model96 = await tf.loadGraphModel("./mobilenet/v2-35_96/model.json");
        console.log("模型 35-96 加载成功");
      } catch (e) {
        console.error("模型 35-96 加载失败", e);
      }

      try {
        model224 = await tf.loadGraphModel("./mobilenet/v2-100_224/model.json");
        console.log("模型 100-224 加载成功");
      } catch (e) {
        console.error("模型 100-224 加载失败", e);
      }
      
      const res = await fetch("./imagenet_labels.json");
      labels = await res.json();      
    }

    async function predict() {
      const imgElement = document.getElementById("inputImage");
      if (!imgElement.src || imgElement.src === "") {
        alert("请先选择一张图片");
        return;
      }

      // 读取图片 → tensor
      //let tensor = tf.browser.fromPixels(imgElement)
      //  .toFloat()
      //  .expandDims(0); // [1, h, w, 3]
        
      let tensor1 = tf.browser.fromPixels(imgElement)
                    .resizeBilinear([96, 96])
                    .toFloat()
                    .expandDims(0)
                    .div(255.0);  

      let tensor2 = tf.browser.fromPixels(imgElement)
                    .resizeBilinear([224, 224])
                    .toFloat()
                    .expandDims(0)
                    .div(255.0);                     
                    

      // 推理 35-96
      if (model96) {
        const logits96 = model96.predict(tensor1);
        const preds96 = await tf.softmax(logits96).data();
        displayResults(preds96, "result96", 5);
      }

      // 推理 100-224
      if (model224) {
        const logits224 = model224.predict(tensor2);
        const preds224 = await tf.softmax(logits224).data();
        displayResults(preds224, "result224", 5);
      }
    }

    function displayResults(preds, textareaId, topK = 5) {
      const textarea = document.getElementById(textareaId);
      // 排序取前 topK
      //let indices = preds.map((p, i) => [p, i]);
      let indices = Array.from(preds, (p, i) => [p, i]);
      indices.sort((a, b) => b[0] - a[0]);
      let top = indices.slice(0, topK);

      let text = top.map(
        //(item, idx) => `${idx + 1}. 类别 ${item[1]} : ${(item[0] * 100).toFixed(2)}%`
        (item, idx) => `${idx + 1}. ${labels[item[1]]} : ${(item[0] * 100).toFixed(2)}%`
      ).join("\n");

      textarea.value = text;
    }

    // 文件上传事件
    document.getElementById("imageUpload").addEventListener("change", function (e) {
      const file = e.target.files[0];
      if (!file) return;
      const reader = new FileReader();
      reader.onload = function (ev) {
        const img = document.getElementById("inputImage");
        img.src = ev.target.result;
        img.onload = function () {
          predict();
        }
      };
      reader.readAsDataURL(file);
    });

    // 页面加载时预加载模型
    loadModels();
  </script>
</body>
</html>
